#!/usr/bin/env python

import sys

__author__ = "Ryan Bressler"
__author__ = "Hector Rovira"

# Establishes a threshold for the distance between Gene TSS and Motif location
maxDistance = 100000

# Merges data from the gene mappings file with the source input file containing chromosome regions
# Inputs
# - Gene mapping file (gene to chromosome locations)
# - Data file (contains chromosome positions and data)
# Outputs
# - Prints mappings from both files within certain distance to standard out
def mergeData(genefile, datafile):
    genesByChr = loadGenesFromBed(genefile)

    inputLines = open(datafile,"r")
    counter = 0
    for inputLine in inputLines:
        vs = inputLine.rstrip().split("\t")
        chr = vs[0]
        counter = counter + 1
        if chr in genesByChr:
            for matchingItem in findMatching(genesByChr[chr], int(vs[1]), int(vs[2])):
                print matchingItem + "\t" + inputLine.rstrip() + "\t" + str(counter)

    inputLines.close()

# Returns array of matching gene mappings to motifs
# Inputs
# - List of genes with start position and strand
# - Motif start location
# - Motif end location
# Outputs
# - Array of matches, each item is a tab-delimited string of Gene information
def findMatching(geneList, motifStart, motifEnd):
    # print "findMatching(" + str(len(geneList)) + "," + str(motifStart) + "," + str(motifEnd) + ")"
    matching = [];
    for geneObj in geneList:
        TSS = geneObj["start"]
        gene_strand = geneObj["strand"]
        gene_id = geneObj["gene_id"]
        gene = geneObj["gene"]

        # Calculate both Upstream and Downstream distance between Gene TSS and Motif
        # - will use Motif START if the Motif is upstream (to-the-left) of Gene TSS
        # - will use Motif END if the Motif is downstream (to-the-right) of Gene TSS
        upstreamDistance = (TSS - motifStart)
        downstreamDistance = (motifEnd - TSS)
        if gene_strand == "-":
            # Flip the calculation if the Gene is on the reverse strand
            upstreamDistance = (TSS - motifEnd)
            downstreamDistance = (motifStart - TSS)

        # TODO: Eventually reduce amount of information being returned (strand and TSS not needed)
        if upstreamDistance > 0 and upstreamDistance <= maxDistance:
            # Motif is upstream of Gene TSS and distance is less/equal than threshold
            # print "findMatching(MATCH_UP," + str(motifStart) + "," + str(motifEnd) + "):[" + gene + "," + str(upstreamDistance) + "," + str(downstreamDistance) + "]"
            matching.append(gene + "\t" + gene_id + "\t" + str(upstreamDistance) + "\t-1\t" + gene_strand + "\t" + str(TSS))
        elif downstreamDistance > 0 and downstreamDistance <= maxDistance:
            # Motif is downstream of Gene TSS and distance is less/equal than threshold
            # print "findMatching(MATCH_DN," + str(motifStart) + "," + str(motifEnd) + "):[" + gene + "," + str(upstreamDistance) + "," + str(downstreamDistance) + "]"
            matching.append(gene + "\t" + gene_id + "\t-1\t" + str(downstreamDistance) + "\t" + gene_strand + "\t" + str(TSS))

    return matching

# Test Function
# Inputs
# - Gene mappings file
# Outputs
# - Prints counts of records by gene id
def showCountsByChromosome(filename):
    genesByChr = loadGenesFromBed(filename)
    for chr in genesByChr:
        print chr + ":" + str(len(genesByChr[chr]))

# Extract Gene Identifiers (plus start and strand) from the given file
# Inputs
# - Gene mappings file
# Outputs
# - List of genes with start position and strand
def loadGenesFromBed(filename):
    colChr = 0;
    colGene = 1;
    colGeneId = 2;
    colStart = 4;
    colStrand = 6;

    genesByChromosome={}
    inputLines = open(filename,"r")
    for inputLine in inputLines:
        vs = inputLine.rstrip().split("\t")
        chr = vs[colChr]
        if not chr in genesByChromosome:
            genesByChromosome[chr]=[]
        genesByChromosome[chr].append({"gene":vs[colGene],"start":int(vs[colStart]),"gene_id":vs[colGeneId],"strand":vs[colStrand]})
    inputLines.close()
    return genesByChromosome

# Main method, executes from command-line
# Inputs
# - Path to Gene mappings file
# - Path to data file with chromosome regions
if __name__ == "__main__":
    if len(sys.argv) > 2:
        #showCountsByChromosome(sys.argv[1])
        #print "############################################################"
        mergeData(sys.argv[1],sys.argv[2])
    else:
        showCountsByChromosome(sys.argv[1])
